from flask import Flask, render_template, request, session
from datetime import datetime

app = Flask(__name__)
app.secret_key = 'your_secret_key'  # 用于 session

def get_life_milestones_v2(birth_date):
    # 只保留指定节点
    milestones = [
        ("出生", 0, birth_date.month),
        ("上幼儿园", 3, 9),
        ("上小学", 6, 9),
        ("上初中", 12, 9),
        ("中考", 15, 6),
        ("高考", 18, 6),
        ("大学本科毕业", 22, 6),
        ("博士毕业", 28, 6),
        ("结婚", 30, 10),
        ("孩子出生", 32, 6),
        ("退休", 60, 1)
    ]
    result = {}
    for name, add_years, fixed_month in milestones:
        year = birth_date.year + add_years
        month = fixed_month
        result[name] = (year, month)
    return result

def calc_age(birth_date, target_year, target_month):
    year = target_year - birth_date.year
    month = target_month - birth_date.month
    if month < 0:
        year -= 1
        month += 12
    if year < 0:
        return None
    return (year, month)

def get_kid_milestones_v2(family, children, child_idx):
    child_role, child_birth = children[child_idx]
    milestones = get_life_milestones_v2(child_birth)
    result = []
    for node, (y, m) in milestones.items():
        dad_age = calc_age(family[0][1], y, m)
        mom_age = calc_age(family[1][1], y, m)
        child_age = calc_age(child_birth, y, m)
        family_ages = [
            f"爸爸: {dad_age[0]}岁{dad_age[1]}月" if dad_age else "爸爸: 未出生",
            f"妈妈: {mom_age[0]}岁{mom_age[1]}月" if mom_age else "妈妈: 未出生"
        ]
        for idx, (krole, kbirth) in enumerate(children):
            if idx == child_idx:
                continue
            age = calc_age(kbirth, y, m)
            if age is None:
                family_ages.append(f"{krole}: 未出生")
            else:
                family_ages.append(f"{krole}: {age[0]}岁{age[1]}月")
        result.append({
            'node': node,
            'date': f"{y}年{m:02d}月",
            'child_age': f"{child_age[0]}岁{child_age[1]}月" if child_age else "未出生",
            'family_ages': '\n'.join(family_ages)
        })
    return child_role, result

def get_all_kids_milestones(family, children):
    all_data = []
    for idx in range(len(children)):
        child_role, milestones = get_kid_milestones_v2(family, children, idx)
        all_data.append({'child_role': child_role, 'milestones': milestones})
    return all_data

@app.route('/', methods=['GET', 'POST'])
def index():
    error = None
    milestones = None
    child_role = None
    all_kids_milestones = None
    # 默认值（爸爸 1985-01-01，妈妈 1987-01-01，第一个孩子 2015-01-01）
    default_dad = session.get('dad_birth', '1985-01-01')
    default_mom = session.get('mom_birth', '1987-01-01')
    default_num_kids = int(session.get('num_kids', 1))
    default_kids = session.get('kids_birth', ['2015-01-01'])
    if request.method == 'POST':
        try:
            dad_birth = request.form['dad_birth']
            mom_birth = request.form['mom_birth']
            num_kids = int(request.form['num_kids'])
            children = []
            kids_birth = []
            for i in range(num_kids):
                kid_birth = request.form[f'kid_birth_{i}']
                kids_birth.append(kid_birth)
                children.append((f"孩子{i+1}", datetime.strptime(kid_birth, "%Y-%m-%d").date()))
            # 存入 session
            session['dad_birth'] = dad_birth
            session['mom_birth'] = mom_birth
            session['num_kids'] = num_kids
            session['kids_birth'] = kids_birth
            family = [("爸爸", datetime.strptime(dad_birth, "%Y-%m-%d").date()), ("妈妈", datetime.strptime(mom_birth, "%Y-%m-%d").date())]
            if num_kids > 1:
                all_kids_milestones = get_all_kids_milestones(family, children)
            else:
                child_role, milestones = get_kid_milestones_v2(family, children, 0)
        except Exception as e:
            error = "输入有误，请检查日期格式和数字。"
    else:
        if isinstance(default_kids, str):
            import ast
            default_kids = ast.literal_eval(default_kids)
    return render_template('index.html', error=error, milestones=milestones, child_role=child_role,
                           default_dad=default_dad, default_mom=default_mom, default_num_kids=default_num_kids,
                           default_kids=default_kids, all_kids_milestones=all_kids_milestones)

if __name__ == '__main__':
    app.run(debug=True) 